#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import glob
import json
import argparse
from pathlib import Path


class LabelConverter:
    """
    LabelConverter
    =============

    功能：
        转换YOLO格式标注文件的标签索引顺序，并自动更新classes.txt和notes.json文件。

    输入：
        - project_path: 项目路径（包含labels/、notes.json）
        - mapping_type: 映射类型（13to7或7to7）

    输出：
        - labels/: 转换后的YOLO格式标注文件
        - classes.txt: 更新后的类别文件
        - notes.json: 更新后的配置文件

    使用方法：
        converter = LabelConverter(project_path, mapping_type)
        converter()
    """

    def __init__(self, project_path, mapping_type):
        self.project_path = Path(project_path)
        self.mapping_type = mapping_type
        self.labels_dir = self.project_path / "labels"
        self.notes_path = self.project_path / "notes.json"
        self.classes_path = self.project_path / "classes.txt"

        # 定义映射关系和新的类别顺序
        self.mappings = {
            '13to7': {
                'mapping': {
                    '0': '5',  # Disease -> 5
                    '1': '5',  # Disease_Ripe -> 5
                    '2': '5',  # Disease_Ripe2 -> 5
                    '3': '5',  # Disease_Ripe4 -> 5
                    '4': '5',  # Disease_Ripe7 -> 5
                    '5': '5',  # Disease_Unripe -> 5
                    '6': '0',  # Ripe -> 0
                    '7': '3',  # Ripe2 -> 3
                    '8': '2',  # Ripe4 -> 2
                    '9': '1',  # Ripe7 -> 1
                    '10': '99',  # Stem -> 删除
                    '11': '4',  # Unripe -> 4
                    '12': '6'  # strawberry_flower -> 6
                },
                'new_classes': [
                    'Ripe',  # 0
                    'Ripe7',  # 1
                    'Ripe4',  # 2
                    'Ripe2',  # 3
                    'Unripe',  # 4
                    'Disease',  # 5
                    'strawberry_flower'  # 6
                ]
            },
            '7to7': {
                'mapping': {
                    '0': '5',
                    '1': '0',
                    '2': '3',
                    '3': '2',
                    '4': '1',
                    '5': '4',
                    '6': '6'
                },
                'new_classes': [
                    'Ripe',  # 0
                    'Ripe7',  # 1
                    'Ripe4',  # 2
                    'Ripe2',  # 3
                    'Unripe',  # 4
                    'Disease',  # 5
                    'strawberry_flower'  # 6
                ]
            }
        }

        if mapping_type not in self.mappings:
            raise ValueError(f"不支持的映射类型: {mapping_type}，可选: 13to7, 7to7")

        self.mapping = self.mappings[mapping_type]['mapping']
        self.new_classes = self.mappings[mapping_type]['new_classes']

    def process_file(self, file_path):
        """处理单个文件，修改每行的第一个数字"""
        with open(file_path, 'r', encoding='utf-8') as file:
            lines = file.readlines()

        modified_lines = []
        for line in lines:
            parts = line.strip().split()
            if not parts:  # 跳过空行
                continue

            # 处理第一个数字
            first_num = parts[0]
            if first_num in self.mapping:
                new_num = self.mapping[first_num]
                # 如果映射为99，则跳过这一行（删除）
                if new_num == '99':
                    continue
                parts[0] = new_num
            # 如果数字不在映射表中，保持不变

            modified_lines.append(' '.join(parts))

        # 写回文件
        with open(file_path, 'w', encoding='utf-8') as file:
            file.write('\n'.join(modified_lines))

    def process_labels(self):
        """处理labels目录中的所有txt文件"""
        if not self.labels_dir.exists():
            print(f"错误: {self.labels_dir} 不存在")
            return

        print(f"处理文件夹: {self.labels_dir}")
        print(f"使用映射表: {self.mapping_type}")

        # 处理该labels目录中的所有txt文件
        txt_files = [f for f in os.listdir(self.labels_dir) if f.endswith('.txt')]

        if not txt_files:
            print("警告: 未找到任何txt文件")
            return

        for filename in txt_files:
            file_path = self.labels_dir / filename
            self.process_file(file_path)
            print(f"已处理: {filename}")

        print(f"完成: 共处理 {len(txt_files)} 个文件")

    def update_classes_txt(self):
        """更新classes.txt文件"""
        try:
            with open(self.classes_path, 'w', encoding='utf-8') as f:
                for class_name in self.new_classes:
                    f.write(f"{class_name}\n")
            print(f"已更新: {self.classes_path}")
        except Exception as e:
            print(f"更新classes.txt失败: {e}")

    def update_notes_json(self):
        """更新notes.json文件"""
        if not self.notes_path.exists():
            print(f"警告: {self.notes_path} 不存在，跳过更新")
            return

        try:
            with open(self.notes_path, 'r', encoding='utf-8') as f:
                notes_data = json.load(f)

            # 更新categories部分
            new_categories = []
            for i, class_name in enumerate(self.new_classes):
                new_categories.append({
                    "id": i,
                    "name": class_name
                })

            notes_data["categories"] = new_categories

            # 写回文件
            with open(self.notes_path, 'w', encoding='utf-8') as f:
                json.dump(notes_data, f, indent=2, ensure_ascii=False)

            print(f"已更新: {self.notes_path}")
        except Exception as e:
            print(f"更新notes.json失败: {e}")

    def __call__(self):
        print(f"开始标签顺序转换 ...")
        print(f"项目路径: {self.project_path}")
        print(f"映射类型: {self.mapping_type}")

        # 1. 处理labels目录中的标注文件
        self.process_labels()

        # 2. 更新classes.txt文件
        self.update_classes_txt()

        # 3. 更新notes.json文件
        self.update_notes_json()

        print("标签顺序转换完成!")


def main():
    parser = argparse.ArgumentParser(description='YOLO标签顺序转换工具')
    parser.add_argument('project_path', type=str, help='项目路径（包含labels/、notes.json）')
    parser.add_argument('mapping_type', choices=['13to7', '7to7'],
                        help='映射类型：13to7（十三分类转七分类）或 7to7（七分类转七分类）')

    args = parser.parse_args()

    converter = LabelConverter(args.project_path, args.mapping_type)
    converter()


if __name__ == "__main__":
    main()